# encoding: utf-8

from model.knn import fit_model_k_fold, predict, fit_model_shuffle
from prepare.load_data import(
    load_source, 
    load_all_table, 
    load_all_table_source, 
    load_item_table, 
    load_local_infects_data, 
    load_local_dead_data,
    load_gpd_source,
    load_area_name,
    load_profession,
    load_profession_infect_data,
    load_profession_dead_data,
)
from visual.visual import imshow
import matplotlib.pyplot as plt
plt.rcParams['font.family'] = ['sans-serif']
plt.rcParams['font.sans-serif'] = ['SimHei']
import config
from model.grayforecast import GrayForecast, grayforecast2
from prepare.lcoal_prob import to_local_a_prob
import numpy as np
import pandas as pd
from model.multipoly_regression import train_multi
# from utils.tasks import all_trainer
from model.knn import fit_model_k_fold
from model.poly import poly
from model.regressor import (rfr, dtr, adbt, br, etr)

print('加载数据文件：')
source = load_source(config.SOURCE_DATA_PATH)
gdp_source = load_gpd_source(config.SOURCE_GDP_DATA_PATH)
print('加载完成...')


####################################################################### 第一题
# 训练总表数据
pd.set_option('display.max_columns', 10000, 'display.max_rows', 10000)
def one_infect():
    # R = GrayForecast(all['发病数.4']).forecast(3, 5)
    # A = GrayForecast(all['发病数']).forecast(3, 7)
    # B = GrayForecast(all['发病数.1']).forecast(3, 3)
    # C = GrayForecast(all['发病数.2']).forecast(3, 3)
    # D = GrayForecast(all['发病数.3']).forecast(3, 4)
    # print(A+B+C+D)
    # print(R)

    all = load_all_table_source(source)
    R = GrayForecast(all['发病数.4']).forecast(3, 3)
    A = GrayForecast(all['发病数']).forecast(3, 3)
    B = GrayForecast(all['发病数.1']).forecast(3, 3)
    C = GrayForecast(all['发病数.2']).forecast(3, 3)
    D = GrayForecast(all['发病数.3']).forecast(3, 3)
    
    r = A+B+C+D
    plt.scatter([2004,2005,2006,2007,2008,2009,2010,2011,2012,2013,2014,2015,2016], all['发病数.4'], c="black", label="当前感染人数")
    plt.plot([2004,2005,2006,2007,2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018,2019], R, c="red", label="总感染人数预测模型")
    plt.plot([2004,2005,2006,2007,2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018,2019], r, c="blue", label="诊断方案分量预测模型")
    plt.title("2004-2019感染人数预测")
    plt.legend(loc='upper right')
    plt.xlabel("年份")
    plt.ylabel("感染人数")
    plt.xticks(np.arange(2004, 2020, 1))
    # plt.show()
    print(r)
    print(R)
    print(f"2019感染人数预测：{R.iloc[-1]}")
    return R,r

def one_dead():
    all = load_all_table_source(source)
    R = GrayForecast(all['发病数.4']).forecast(3, 3)
    A = GrayForecast(all['发病数']).forecast(3, 3)
    B = GrayForecast(all['发病数.1']).forecast(3, 3)
    C = GrayForecast(all['发病数.2']).forecast(3, 3)
    D = GrayForecast(all['发病数.3']).forecast(3, 3)
    Dr = all['死亡数.4']

    # R = (R+(A+B+C+D))/2

    R1 = R[:-3]
    R2 = R[-3:]

    # t,tr = train_multi(R1.values, Dr)
    # res = t.predict(tr.fit_transform(R2.values))
    t = poly(R1.values, Dr.values)
    res = t.predict(R2.values)
    res = np.append(Dr, res)
    plt.scatter([2004,2005,2006,2007,2008,2009,2010,2011,2012,2013,2014,2015,2016], all['死亡数.4'], c="black", label="当前死亡数")
    plt.plot([2004,2005,2006,2007,2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018,2019], res, c="blue", label="预测模型")
    plt.title("2004-2019死亡人数预测")
    plt.legend(loc="upper right")
    plt.xlabel("年份")
    plt.ylabel("死亡人数")
    plt.xticks(np.arange(2004, 2020, 1))
    print("死亡人数：")
    print(res)
    print(f"2019预测死亡人数{res[-1]}")
    plt.show()
    

# one_infect()
one_dead()

###################################################################### 第二题
def two_anwser_2_infect_areas():
    areas = load_area_name(source)
    item_infect_all = load_local_infects_data(source)
    item_dead_all = load_local_dead_data(source)
    counts, k = item_infect_all.shape
    inv_infect_all = item_infect_all.T.reset_index(drop=True)
    inv_dead_all = item_dead_all.T.reset_index(drop=True)

    # for i in range(31):
    #     plt.scatter(inv_infect_all.iloc[:, i], inv_dead_all.iloc[:, i])
    #     plt.xlabel(areas[i])
    #     plt.show()
    
    infect_res_2019 = []
    for i in range(counts):
        a = GrayForecast(inv_infect_all.iloc[:, i]).forecast(1)        
        infect_res_2019.append(int(a.iloc[0][0]))
    infect_res_2019 = np.array(infect_res_2019)
    top = np.argsort(infect_res_2019)[::-1][:15]
    print("2019年感染重灾区是:")
    print(areas.iloc[top])
    print(infect_res_2019[top])

    dead_res_2019 = []
    for i in range(counts):
        a = GrayForecast(inv_dead_all.iloc[:, i]).forecast(1)        
        dead_res_2019.append(int(a.iloc[0][0]))
    dead_res_2019 = np.array(dead_res_2019)
    dtop = np.argsort(dead_res_2019)[::-1][:15]
    print("2019年死亡重灾区是:")
    print(areas.iloc[dtop])
    print(dead_res_2019[dtop])


def two_answer_2_profession():
    areas = load_profession(source)
    item_infect_all = load_profession_infect_data(source)
    item_dead_all = load_profession_dead_data(source)
    counts, k = item_infect_all.shape
    inv_infect_all = item_infect_all.T.reset_index(drop=True)
    inv_dead_all = item_dead_all.T.reset_index(drop=True)
    inv_dead_all += 1

    # for i in range(31):
    #     plt.scatter(inv_infect_all.iloc[:, i], inv_dead_all.iloc[:, i])
    #     plt.xlabel(areas[i])
    #     plt.show()
    
    infect_res_2019 = []
    for i in range(counts):
        a = GrayForecast(inv_infect_all.iloc[:, i]).forecast(1)        
        infect_res_2019.append(int(a.iloc[0][0]))
    infect_res_2019 = np.array(infect_res_2019)
    top = np.argsort(infect_res_2019)[::-1][:15]
    print("2019年感染职业人群是:")
    print(areas.iloc[top])
    print(infect_res_2019[top])

    dead_res_2019 = []
    for i in range(counts):
        a = GrayForecast(inv_dead_all.iloc[:, i]).forecast(1)        
        dead_res_2019.append(int(a.iloc[0][0]))
    dead_res_2019 = np.array(dead_res_2019)
    dtop = np.argsort(dead_res_2019)[::-1][:15]
    print("2019年死亡职业人群是:")
    print(areas.iloc[dtop])
    print(dead_res_2019[dtop])

# two_answer_2_profession()
# two_anwser_2_infect_areas()

##################################################################### 第三题
import numpy as np

def three():
    gdps = gdp_source.iloc[:, 1:]
    item_all = load_local_infects_data(source)

    print(gdps)
    print(item_all)
    for i in range(5):
        plt.scatter(gdps.iloc[:, i]/1000, item_all.iloc[:, i]/100, c='blue', vmin=120000)
        plt.xlabel('人均gdp')
        plt.ylabel('感染人数')
        plt.xticks(np.arange(0, 100, 10))
        plt.yticks(np.arange(0, 1000, 100))
        plt.show()

